Memory Management

from typing import TypedDict, Annotated
from langgraph.checkpoint.memory import MemorySaver
from langgraph.constants import START, END
from langgraph.graph import StateGraph
from langgraph.graph.message import add_messages
from langchain_ollama import ChatOllama

class State(TypedDict):
    """
    定义图结构中节点间传递的状态结构

    Attributes:
        messages: 消息列表，使用add_messages函数进行合并
    """
    messages: Annotated[list, add_messages]

# 创建状态图构建器
graph_builder = StateGraph(State)

# 初始化本地大语言模型，配置基础URL、模型名称和推理模式
llm = ChatOllama(base_url="http://localhost:11434", model="qwen3:14b", reasoning=False)

def chatbot(state: State):
    """
    聊天机器人节点函数，处理输入消息并生成回复

    Args:
        state (State): 包含消息历史的状态字典

    Returns:
        dict: 包含新生成消息的字典，格式为{"messages": [回复消息]}
    """
    return {"messages": [llm.invoke(state["messages"])]}

# 将聊天机器人节点添加到图中
graph_builder.add_node("chatbot", chatbot)

# 添加从开始节点到聊天机器人节点的边
graph_builder.add_edge(START, "chatbot")

# 添加从聊天机器人节点到结束节点的边
graph_builder.add_edge("chatbot", END)

# 创建内存保存器用于保存对话状态
memory = MemorySaver()

# 编译图结构并设置检查点保存器
graph = graph_builder.compile(checkpointer=memory)

# 绘制图结构并保存为PNG图片
graph.get_graph().draw_png('./graph.png')

# 配置对话线程ID
config = {"configurable": {"thread_id": "chat-1"}}

# 第一次对话：发送初始消息
msg1 = graph.invoke({"messages": ["你好，我叫崔亮，喜欢学习。"]}, config=config)
msg1["messages"][-1].pretty_print()

# 第二次对话：基于上下文询问用户信息
msg2 = graph.invoke({"messages": ["我叫什么？我喜欢做什么？"]}, config=config)
msg2["messages"][-1].pretty_print()

import uuid
from typing import TypedDict, Annotated
import dotenv
from langchain_ollama import ChatOllama
from langchain_core.runnables import RunnableConfig
from langgraph.constants import END, START
from langgraph.graph import StateGraph, MessagesState, add_messages
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.store.memory import InMemoryStore
from langgraph.store.base import BaseStore

# 加载环境变量配置
dotenv.load_dotenv()
# 初始化本地大语言模型，配置模型名称和推理模式
model = ChatOllama(model="qwen3:14b", reasoning=False)

class State(TypedDict):
    """
    定义图中状态的数据结构。

    属性:
        messages (Annotated[list, add_messages]): 使用 add_messages 合并的消息列表。
    """
    messages: Annotated[list, add_messages]

def save_memory(store: BaseStore, user_id: str, content: str):
    """
    将用户输入的内容保存为记忆。

    参数:
        store (BaseStore): 存储系统的实例，用于持久化数据。
        user_id (str): 用户唯一标识符。
        content (str): 需要存储的文本内容。
    """
    namespace = ("memories", user_id)
    store.put(namespace, str(uuid.uuid4()), {"data": content})

def recall_memories(store: BaseStore, user_id: str, query: str, limit: int = 5):
    """
    根据查询语句检索与用户相关的记忆。

    参数:
        store (BaseStore): 存储系统的实例。
        user_id (str): 用户唯一标识符。
        query (str): 查询关键词或句子。
        limit (int, optional): 返回的记忆条数上限，默认是 5 条。

    返回:
        list[str]: 匹配的记忆内容列表。
    """
    namespace = ("memories", user_id)
    memories = store.search(namespace, query=query, limit=limit)
    return [m.value["data"] for m in memories]

def chatbot(state: MessagesState, config: RunnableConfig, *, store: BaseStore):
    """
    聊天机器人主逻辑节点函数。

    参数:
        state (MessagesState): 当前对话的状态信息，包括历史消息等。
        config (RunnableConfig): 运行时配置信息，如线程ID、用户ID等。
        store (BaseStore): 用于读取和写入用户记忆的存储接口。

    返回:
        dict: 更新后的消息状态字典。
    """
    user_id = config["configurable"]["user_id"]

    # 检索历史记忆
    query = state["messages"][-1].content
    related_memories = recall_memories(store, user_id, query)

    # 构造系统提示
    system_msg = (
        "你是一个友好的聊天助手。\n"
        f"以下是关于用户的记忆:\n{chr(10).join(related_memories) if related_memories else '暂无'}"
    )

    # 保存当前消息到记忆
    save_memory(store, user_id, query)

    # 调用模型生成回复
    response = model.invoke(
        [{"role": "system", "content": system_msg}] + state["messages"]
    )
    return {"messages": response}

# 创建状态图并定义流程
builder = StateGraph(State)
builder.add_node(chatbot)
builder.add_edge(START, "chatbot")
builder.add_edge("chatbot", END)

# 初始化检查点和存储组件
checkpointer = InMemorySaver()
store = InMemoryStore()

# 编译构建最终可运行的图对象，并绘制其结构图
graph = builder.compile(
    checkpointer=checkpointer,
    store=store,
)
graph.get_graph().draw_png('./graph.png')

# 第一次交互测试：记录用户基本信息
config1 = {"configurable": {"thread_id": "1", "user_id": "1"}}
msg1 = graph.invoke({"messages": [{"role": "user", "content": "我叫崔亮，喜欢学习。"}]}, config1)
print("第一次回复：")
msg1["messages"][-1].pretty_print()

# 第二次交互测试：验证是否能回忆起之前的信息
config2 = {"configurable": {"thread_id": "2", "user_id": "1"}}
msg2 = graph.invoke({"messages": [{"role": "user", "content": "我叫什么？我喜欢做什么？"}]}, config2)
print("第二次回复：")
msg2["messages"][-1].pretty_print()

import dotenv
from langchain_core.messages.utils import trim_messages, count_tokens_approximately
from langchain_ollama import ChatOllama
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.prebuilt import create_react_agent

# 加载环境变量配置
dotenv.load_dotenv()
# 初始化本地大语言模型，配置模型名称和推理模式
model = ChatOllama(model="qwen3:14b", reasoning=False)
# 定义工具列表，
tools = []

def pre_model_hook(state):
    """
    在模型处理前对消息进行预处理的钩子函数

    该函数用于裁剪消息历史，只保留最近的若干条消息，避免上下文过长

    Args:
        state (dict): 包含对话状态的字典，其中"messages"键对应消息列表

    Returns:
        dict: 包含裁剪后消息的字典，键为"llm_input_messages"
    """
    # 参数说明:
    #   state["messages"]: 需要裁剪的消息列表
    #   strategy: 裁剪策略，"last"表示从最后开始裁剪
    #   token_counter: 用于计算token数量的函数，这里使用近似计算方法
    #   max_tokens: 最大token数量限制，设置为300
    #   start_on: 开始裁剪的消息类型，"human"表示从人类用户的消息开始
    #   end_on: 结束裁剪的消息类型，可以是"human"或"tool"类型的消息
    # 返回值: 裁剪后的消息列表
    trimmed_messages = trim_messages(
        state["messages"],
        strategy="last",
        token_counter=count_tokens_approximately,
        max_tokens=300,
        start_on="human",
        end_on=("human", "tool"),
    )

    return {"llm_input_messages": trimmed_messages}

checkpointer = InMemorySaver()
agent = create_react_agent(
    model,
    tools,
    pre_model_hook=pre_model_hook,
    checkpointer=checkpointer,
)
config = {"configurable": {"thread_id": "user-001"}}
msg1 = agent.invoke({"messages": [("user", "你好，我叫崔亮")]}, config)
msg1["messages"][-1].pretty_print()
like_list = ['唱', '跳', 'rap', '篮球']
for i in like_list:
    msg = "我喜欢做的事是：" + i
    print(msg)
    agent.invoke({"messages": [("user", msg)]}, config)
msg2 = agent.invoke({"messages": [("user", "我叫什么？我喜欢做的事是什么？")]}, config)
msg2["messages"][-1].pretty_print()

import dotenv
from langchain_core.messages.utils import trim_messages, count_tokens_approximately
from langchain_ollama import ChatOllama
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.prebuilt import create_react_agent
from langgraph.prebuilt.chat_agent_executor import AgentState
from langmem.short_term import SummarizationNode, RunningSummary

# 加载环境变量配置
dotenv.load_dotenv()

# 初始化本地大语言模型，配置模型名称和推理模式
model = ChatOllama(model="qwen3:14b", reasoning=False)

# 定义工具列表，
tools = []

# 创建一个SummarizationNode实例，用于处理文本摘要任务
# 参数说明:
#   token_counter: 用于估算文本token数量的函数，这里使用count_tokens_approximately函数
#   model: 指定使用的语言模型实例
#   max_tokens: 限制处理文本的最大token数量为300
#   max_summary_tokens: 限制生成摘要的最大token数量为128
#   output_messages_key: 指定输出消息在结果中的键名，设置为"llm_input_messages"
summarization_node = SummarizationNode(
    token_counter=count_tokens_approximately,
    model=model,
    max_tokens=300,
    max_summary_tokens=128,
    output_messages_key="llm_input_messages",
)

# 自定义状态类，继承自AgentState，添加上下文字段用于存储运行时摘要信息
class State(AgentState):
    context: dict[str, RunningSummary]

# 初始化内存检查点保存器，用于持久化代理状态
checkpointer = InMemorySaver()

# 创建React代理，整合模型、工具、摘要节点和状态管理器
agent = create_react_agent(
    model=model,
    tools=tools,
    pre_model_hook=summarization_node,
    state_schema=State,
    checkpointer=checkpointer,
)

# 配置线程ID，用于标识用户会话
config = {"configurable": {"thread_id": "user-001"}}

# 启动对话，发送用户自我介绍消息并获取模型响应
msg1 = agent.invoke({"messages": [("user", "你好，我叫崔亮")]}, config)
msg1["messages"][-1].pretty_print()

# 定义用户兴趣列表
like_list = ['唱', '跳', 'rap', '篮球']

# 循环发送用户兴趣信息，逐条更新上下文
for i in like_list:
    msg = "我喜欢做的事是：" + i
    print(msg)
    agent.invoke({"messages": [("user", msg)]}, config)

# 查询用户姓名和兴趣，测试模型对上下文的理解能力
msg2 = agent.invoke({"messages": [("user", "我叫什么？我喜欢做的事是什么？")]}, config)
msg2["messages"][-1].pretty_print()